﻿using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Forms;
using Ast.Core;
using Ast.Core.Events;
using Ast.Core.Services;
using Ast.Framework.Forms;
using Castle.Core.Logging;
using Equin.ApplicationFramework;

namespace Ast
{
    [FormMetadata(FormType.Hook, IsTabForm = true)]
    public partial class HookForm : Form, IKernelReloadEvents, IHookSystemEvents
    {
        private readonly ISsdtService _ssdtService;
        private readonly IHookSystemService _hookSystemService;
        private readonly IKernelReloadService _kernelReloadService;
        private readonly IDeviceController _deviceController;

        private const string MonitorInfoEventName = "Global\\MonitorInfoEvent";
        private List<FunctionViewItem> _functions;
        private BindingListView<FunctionViewItem> _functionListView;
        private uint _reloadedKernelOffset;

        private readonly object _syncRoot = new object();

        public HookForm(ISsdtService ssdtService, IHookSystemService hookSystemService, IKernelReloadService kernelReloadService, IDeviceController deviceController)
        {
            _ssdtService = ssdtService;
            _hookSystemService = hookSystemService;
            _kernelReloadService = kernelReloadService;
            _deviceController = deviceController;

            InitializeComponent();

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        private void HookForm_Load(object sender, System.EventArgs e)
        {
            // disable button until system is hooked
            btnMonitor.Enabled = false;
            btnRemove.Enabled = false;

            btnReloadFunctions.PerformClick();
        }

        private async void btnReloadFunctions_Click(object sender, EventArgs e)
        {
            await Task.Run(() =>
            {
                lock (_syncRoot)
                {
                    if (!_ssdtService.FindAddress()) return;

                    // get items
                    _functions = new List<FunctionViewItem>();
                    _ssdtService.ListFunctions(_ssdtService.SsdtShadow, function =>
                    {
                        _functions.Add(new FunctionViewItem
                        {
                            IsShadowFunction = !function.Module.Name.StartsWith("nt"),
                            ServiceTableBase = function.ServiceTableBase,
                            Index = function.FunctionIndex,
                            Name = function.FunctionName,
                            Address = function.FunctionAddress
                        });
                        return true;
                    });

                    // set view data source
                    lbFunctions.BeginUpdate();
                    _functionListView = new BindingListView<FunctionViewItem>(_functions);
                    ApplyViewFilter();
                    lbFunctions.DataSource = _functionListView;
                    lbFunctions.EndUpdate();
                }
            });
        }

        private void txbFilter_TextChanged(object sender, EventArgs e)
        {
            ApplyViewFilter();
        }

        private void ckbShowShadowSsdt_CheckedChanged(object sender, EventArgs e)
        {
            ckbUseReloadedFunction.Visible = !ckbShowShadowSsdt.Checked;
            ApplyViewFilter();
        }

        private void ApplyViewFilter()
        {
            if (_functionListView == null) return;

            if (string.IsNullOrEmpty(txbFilter.Text))
            {
                _functionListView.ApplyFilter(item => item.IsShadowFunction == ckbShowShadowSsdt.Checked);
            }
            else
            {
                int index;
                if (int.TryParse(txbFilter.Text, NumberStyles.HexNumber, CultureInfo.CurrentCulture, out index))
                {
                    _functionListView.ApplyFilter(item => item.IsShadowFunction == ckbShowShadowSsdt.Checked && item.Index >= index); // by index
                }
                else
                {
                    _functionListView.ApplyFilter(item => item.IsShadowFunction == ckbShowShadowSsdt.Checked && item.Name.IndexOf(txbFilter.Text, StringComparison.InvariantCultureIgnoreCase) >= 0); // by name
                }
            }
        }

        private void btnMonitor_Click(object sender, System.EventArgs e)
        {
            var functionItemView = lbFunctions.SelectedItem as ObjectView<FunctionViewItem>;
            if (functionItemView == null) return;

            var functionItem = functionItemView.Object;
            UpdateMonitorFunctions(functionItem, true);
        }

        private void btnRemove_Click(object sender, EventArgs e)
        {
            var functionItem = lbMonitorFunctions.SelectedItem as FunctionViewItem;
            if (functionItem == null) return;

            UpdateMonitorFunctions(functionItem, false);
        }

        private void UpdateMonitorFunctions(FunctionViewItem functionItem, bool isMonitor)
        {
            // request
            var request = new IoCtl.HookSetMonitorItemRequest
            {
                ServiceTableBase = functionItem.ServiceTableBase,
                FunctionIndex = functionItem.Index,
                IsMonitor = isMonitor
            };
            if (ckbUseReloadedFunction.Checked && !functionItem.IsShadowFunction)
            {
                request.ReloadedFunctionAddress = functionItem.Address + _reloadedKernelOffset;
            }

            // send request
            SetMonitorInfoResponse setMonitorInfoResponse;
            _deviceController.Send(IoCtl.IoCtlHookSetMonitorItem, request, out setMonitorInfoResponse);

            // update monitor functions
            lbMonitorFunctions.BeginUpdate();
            lbMonitorFunctions.Items.Clear();
            foreach (var functionIndex in setMonitorInfoResponse.FunctionIndexes.TakeWhile(functionIndex => functionIndex != -1))
            {
                lbMonitorFunctions.Items.Add(_functions.First(x => x.Index == functionIndex));
            }
            lbMonitorFunctions.EndUpdate();
        }

        private async void btnListen_Click(object sender, EventArgs e)
        {
            EventWaitHandle monitorInfoEvent;
            if (!EventWaitHandle.TryOpenExisting(MonitorInfoEventName, EventWaitHandleRights.Synchronize, out monitorInfoEvent))
            {
                Logger.Win32ErrorFormat("failed to open monitor info event");
                return;
            }

            btnListen.Enabled = false;
            btnListen.Text = "Listening";

            await Task.Run(() =>
            {
                while (monitorInfoEvent.WaitOne())
                {
                    GetMonitorInfoResponse getMonitorInfoResponse;
                    if (!_deviceController.Send(IoCtl.IoCtlHookGetMonitorInfo, out getMonitorInfoResponse))
                    {
                        Logger.Error("failed to get monitor info");
                        return;
                    }

                    foreach (var monitorInfoItem in getMonitorInfoResponse.MonitorInfoItemsDatas)
                    {
                        if (monitorInfoItem.FunctionIndex == -1) break; // no more

                        var funcitonItem = _functions.First(x => x.Index == monitorInfoItem.FunctionIndex);
                        var parametersInfo = new StringBuilder();
                        for (var i = 0; i < monitorInfoItem.ParameterCount; i++)
                        {
                            parametersInfo.AppendFormat(", p{0} 0x{1:x}", i, monitorInfoItem.Parameters[i].ToInt32());
                        }
                        if (monitorInfoItem.ParameterCount > 0)
                        {
                            parametersInfo.Remove(0, 2);
                        }
                        Logger.InfoFormat("pid {0} called function {1} 0x{2:x}, {3}", monitorInfoItem.ProcessId, funcitonItem.Name, monitorInfoItem.FunctionAddress, parametersInfo);
                    }
                }
            }).ConfigureAwait(false);
        }

        private void RemoveAllMonitorFunctions()
        {
            foreach (var functionItem in lbMonitorFunctions.Items.Cast<FunctionViewItem>())
            {
                var request = new IoCtl.HookSetMonitorItemRequest
                {
                    ServiceTableBase = functionItem.ServiceTableBase,
                    FunctionIndex = functionItem.Index,
                    IsMonitor = false
                };
                _deviceController.Send(IoCtl.IoCtlHookSetMonitorItem, request);
            }
            lbMonitorFunctions.Items.Clear();
        }

        private async void btnHook_Click(object sender, EventArgs e)
        {
            await Task.Run(() =>
            {
                _hookSystemService.AcquireLock();
                _hookSystemService.Hook();
                _hookSystemService.ReleaseLock();
            }).ConfigureAwait(false);
        }

        private async void btnReloadKernel_Click(object sender, EventArgs e)
        {
            await Task.Run(() =>
            {
                _kernelReloadService.AcquireLock();
                _kernelReloadService.Reload();
                _kernelReloadService.ReleaseLock();
            }).ConfigureAwait(false);
        }

        #region Events

        void IKernelReloadEvents.OnRealoding()
        {
        }

        void IKernelReloadEvents.OnReloaded(IntPtr curKernelBase, IntPtr newKernelBase)
        {
            Invoke(new Action(() =>
            {
                _reloadedKernelOffset = (uint) (newKernelBase.ToInt32() - curKernelBase.ToInt32());
                ckbUseReloadedFunction.Enabled = true;
            }));
        }

        void IKernelReloadEvents.OnUnloading(IntPtr kernelBase)
        {
            Invoke(new Action(() =>
            {
                RemoveAllMonitorFunctions();

                _reloadedKernelOffset = 0;
                ckbUseReloadedFunction.Checked = false;
                ckbUseReloadedFunction.Enabled = false;
            }));
        }

        void IKernelReloadEvents.OnUnloaded()
        {
        }

        void IHookSystemEvents.OnHooked()
        {
            Invoke(new Action(() =>
            {
                btnMonitor.Enabled = true;
                btnRemove.Enabled = true;
            }));
        }

        void IHookSystemEvents.OnUnhooked()
        {
            Invoke(new Action(() =>
            {
                btnMonitor.Enabled = false;
                btnRemove.Enabled = false;
            }));
        }

        #endregion
    }

    public class FunctionViewItem
    {
        public bool IsShadowFunction { get; set; }
        public uint ServiceTableBase { get; set; }
        public int Index { get; set; }
        public string Name { get; set; }
        public uint Address { get; set; }

        public override string ToString()
        {
            return string.Format("0x{0:x} {1}", Index, Name);
        }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct SetMonitorInfoResponse
    {
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 512)]
        public int[] FunctionIndexes;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct GetMonitorInfoResponse
    {
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 10)]
        public MonitorInfoItemData[] MonitorInfoItemsDatas;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct MonitorInfoItemData
    {
        public uint ProcessId;

        public int FunctionIndex;

        public uint FunctionAddress;

        public int ParameterCount;

        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]
        public IntPtr[] Parameters;
    }
}